from tqdm import tqdm
import pandas as pd
import spider as msd


def main(start_date, end_date, searchKey, migrationType, o_level, d_level):
    """

    Args:
        start_date (str): ''
        end_date (str): ''
        searchKey (str): "地级市名称/省名称"
        migrationType (str): 'in'/'out'
        o_level (str): 出发地行政级别,'city/province'
        d_level (str): 到达地行政级别,'city/province'
    """
    city = pd.read_json('CityCode.json')
    NameList = list(city[searchKey].unique())
    start_date = start_date
    end_date = end_date
    dateRange = pd.date_range(
        start_date, end_date).strftime("%Y-%m-%d").tolist()
    batchdata = []
    with tqdm(total=len(NameList)) as pbar:
        for i in NameList:
            a = msd.GetBatchDateData(o_level, d_level, i, migrationType=migrationType,
                                     date=dateRange)
            batchdata.append(a)
            pbar.update(1)
    writer = pd.ExcelWriter(start_date+'to'+end_date +
                            o_level+'to'+d_level+'.xlsx')
    for j in batchdata:
        j.to_excel(writer, j.iloc[0][2])
    writer.save()


if __name__ == '__main__':
    main('2020-02-01', '2020-03-01', '省名称', 'in', 'province', 'province')
